import json
import os
import random
import numpy as np
import open3d as o3d
from plyfile import PlyData, PlyElement
from PointCloudRegistration import registration_point_clouds

def read_ply(ply_path):
    # 读取header部分（使用二进制模式解决编码问题）
    header = []
    with open(ply_path, 'rb') as f:  # 修改为二进制模式
        for line_bytes in f:
            line = line_bytes.decode('utf-8').strip()  # 显式解码
            header.append(line)
            if line == 'end_header':
                break
    
    # 读取顶点数据（使用plyfile库原生方法）
    ply_data = PlyData.read(ply_path)
    vertex_element = ply_data['vertex']
    vertex_data = []
    property_names = [prop.name for prop in vertex_element.properties]
    
    # 转换每个顶点的属性为字符串列表
    for v in vertex_element:
        row = []
        for prop in vertex_element.properties:
            val = v[prop.name]
            # 根据数据类型格式化字符串
            if 'float' in prop.val_dtype:
                formatted = f"{val:.6f}".rstrip('0').rstrip('.') if '.' in f"{val:.6f}" else f"{val}"
            elif 'int' in prop.val_dtype or 'uint' in prop.val_dtype:
                formatted = f"{val}"
            else:
                formatted = str(val)
            row.append(formatted)
        vertex_data.append(row)
    
    return header, vertex_data, property_names

def write_ply(output_path, header, vertices):
    # 更新顶点数量
    for i, line in enumerate(header):
        if line.startswith('element vertex'):
            header[i] = f'element vertex {len(vertices)}'
            break
    
    with open(output_path, 'w', encoding='utf-8') as f:  # 显式指定编码
        f.write('\n'.join(header) + '\n')
        for v in vertices:
            f.write(' '.join(v) + '\n')

def merge_point_clouds(metadata_path, output_path):
    with open(metadata_path, 'r') as f:
        metadata = json.load(f)
    
    all_vertices = []
    common_header = [
        'ply',
        'format ascii 1.0',
        'element vertex 0',
        'property float x',
        'property float y',
        'property float z',
        'property uchar red',
        'property uchar green',
        'property uchar blue',
        'end_header'
    ]
    
    for group_idx, group in enumerate(metadata["groups"]):
        obj_path = group["output_file"]
        ply_path = os.path.splitext(obj_path)[0] + ".ply"
        translate = group["transform"]["translate"]
        tx, ty, tz = translate
        
        if not os.path.exists(obj_path) or not os.path.exists(ply_path):
            print(f"Warning: {obj_path} or {ply_path} not found. Skipping.")
            continue
        
        # 生成当前组的随机颜色
        group_color = (
            random.randint(0, 255),
            random.randint(0, 255),
            random.randint(0, 255)
        )
        
        # 读取并采样OBJ模型
        try:
            mesh = o3d.io.read_triangle_mesh(obj_path)
        except Exception as e:
            print(f"Error reading {obj_path}: {str(e)}")
            continue
        
        if not mesh.has_vertices():
            print(f"Warning: {obj_path} has no vertices. Skipping.")
            continue
        if not mesh.triangles:
            print(f"Warning: {obj_path} has no surface geometry. Skipping.")
            continue
        
        try:
            target_pcd = mesh.sample_points_uniformly(number_of_points=20000)
        except Exception as e:
            print(f"Sampling failed for {obj_path}: {str(e)}")
            continue
        
        # 读取PLY数据
        try:
            header, vertex_data, property_names = read_ply(ply_path)
        except Exception as e:
            print(f"Error reading {ply_path}: {str(e)}")
            continue
        
        if not vertex_data:
            print(f"Warning: {ply_path} has no vertices. Skipping.")
            continue
        
        # 确定x, y, z属性的索引
        try:
            x_idx = property_names.index('x')
            y_idx = property_names.index('y')
            z_idx = property_names.index('z')
        except ValueError:
            print(f"PLY file {ply_path} missing x, y, or z properties. Skipping.")
            continue
        
        source_points = []
        valid_vertex_indices = []
        for i, v in enumerate(vertex_data):
            try:
                x = float(v[x_idx])
                y = float(v[y_idx])
                z = float(v[z_idx])
                source_points.append([x, y, z])
                valid_vertex_indices.append(i)
            except (ValueError, IndexError):
                continue
        
        if not source_points:
            print(f"Warning: {ply_path} has no valid vertices. Skipping.")
            continue
        
        source_pcd = o3d.geometry.PointCloud()
        source_pcd.points = o3d.utility.Vector3dVector(np.array(source_points))
        
        # ICP配准
        reg_result = registration_point_clouds(ply_path, obj_path)
        
        T = reg_result
        transformed_vertices = []
        for i in valid_vertex_indices:
            v = vertex_data[i]
            try:
                x = float(v[x_idx])
                y = float(v[y_idx])
                z = float(v[z_idx])
            except (ValueError, IndexError):
                continue
            
            p_homog = np.array([x, y, z, 1.0])
            p_transformed = T @ p_homog
            
            x_t = p_transformed[0] + tx
            y_t = p_transformed[1] + ty
            z_t = p_transformed[2] + tz
            
            # 构造新的顶点数据，包含坐标和颜色
            new_v = [
                f"{x_t:.6f}",
                f"{y_t:.6f}",
                f"{z_t:.6f}",
                str(group_color[0]),
                str(group_color[1]),
                str(group_color[2])
            ]
            transformed_vertices.append(new_v)
        
        all_vertices.extend(transformed_vertices)
    
    if len(all_vertices) == 0:
        print("Error: No valid data processed.")
        return
    
    # 更新总顶点数
    for idx, line in enumerate(common_header):
        if line.startswith('element vertex'):
            common_header[idx] = f'element vertex {len(all_vertices)}'
            break
    
    # 写入最终文件
    write_ply(output_path, common_header, all_vertices)
    print(f"Merged point cloud saved to {output_path}")

if __name__ == "__main__":
    merge_point_clouds("output_metadata.json", "merged.ply")